import xml.etree.ElementTree as ET
from tqdm import tqdm
from pathlib import Path
import os

def convert_label(image_path, lb_path):
    names = {0:'person', 1: 'head_front', 2:'head_back', 3: 'head_other'}
    def convert_box(size, box):
        dw, dh = 1. / size[0], 1. / size[1]
        x, y, w, h = (box[0] + box[1]) / 2.0 - 1, (box[2] + box[3]) / 2.0 - 1, box[1] - box[0], box[3] - box[2]
        return x * dw, y * dh, w * dw, h * dh

    xml_path = image_path.with_suffix('.xml')  # corresponding XML file path
    out_file = open(lb_path, 'w')
    
    tree = ET.parse(xml_path)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)
   


    names = list(names.values())  # names list
    for obj in root.iter('object'):
        cls = obj.find('name').text
        # Find the 'attributes' element
        attributes_element = obj.find('.//attributes')

        # Find the 'attribute' element with 'name' equal to 'toward'
        toward_element = attributes_element.find('.//attribute[name="toward"]')

        # Extract the value of 'toward'
        if toward_element is not None:
            toward_value = toward_element.find('value').text
            cls = cls + '_' + toward_value
        
        if cls in names and int(obj.find('difficult').text) != 1:
            xmlbox = obj.find('bndbox')
            bb = convert_box((w, h), [float(xmlbox.find(x).text) for x in ('xmin', 'xmax', 'ymin', 'ymax')])
            cls_id = names.index(cls)  # class id
            out_file.write(" ".join(str(a) for a in (cls_id, *bb)) + '\n')

            

# Set the directory containing image and XML files
if os.environ.get('face_dev'):
    filedirs = ['/root/code/cvmark/face/data/v7', '/root/code/cvmark/face/data/v6']
else:
    filedirs = ['/home/data/2795', '/home/data/2796']
# filedir = '/home/data/2795'
for filedir in filedirs:
    data_dir = Path(filedir)

    # Create output directory for YOLO format labels
    output_dir = Path(filedir)
    output_dir.mkdir(exist_ok=True, parents=True)

    # Get all image files in the directory
    image_files = list(data_dir.glob('*.jpg'))

    for image_path in tqdm(image_files, desc='Converting to YOLO format'):
        lb_path = output_dir / (image_path.stem + '.txt')
        convert_label(image_path, lb_path)